package itsecurity.cert;

import iaik.security.provider.IAIK;
import iaik.utils.Util;
import iaik.x509.X509Certificate;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;
import java.util.Arrays;
import java.util.Enumeration;

import java.security.cert.Certificate;

/**
 * Manages all key-related operations. Low cohesion danger here... 
 *
 */
public class KeyManager {

	/**
	 * KeyStore filename
	 */
	public static final String KEYSTORE = "keyStore";
	
	/**
	 * Keystore pwd
	 */
	private static final String STORE_PWD = "Bob";
//	private PrivateKey privateKey;
//	private X509Certificate[] certs;
	
	/**
	 * Creates a keystore. Password is determined in a static variable, which
	 * is a quick but terrible solution.
	 * 
	 * @param fos determine the location of the keystore beforehand
	 * @throws KeyStoreException
	 * @throws NoSuchProviderException
	 * @throws NoSuchAlgorithmException
	 * @throws CertificateException
	 * @throws IOException
	 */
	public static KeyStore createKeyStore(FileOutputStream fos) throws KeyStoreException, 
	NoSuchProviderException, NoSuchAlgorithmException, CertificateException, IOException {
		KeyStore store = KeyStore.getInstance("IAIKKeyStore", "IAIK");
	  store.load(null, null);
	  store.store(fos, STORE_PWD.toCharArray());
//	  Arrays.fill(password, '\u0000');
	  return store;
	}
	
	/**
	 * 
	 * @param alias
	 * @param password
	 * @param storeFile
	 * @return First value in the certificate chain for alias. This should be changed.
	 */
	public static Certificate getCert(String alias, char[] password, File storeFile) {
		FileInputStream fis;
		KeyStore store=null;
		Certificate cert=null;
		try {
			store = KeyStore.getInstance("IAIKKeyStore", "IAIK");
			fis = new FileInputStream(storeFile);
			store.load(fis, STORE_PWD.toCharArray());
	    cert = store.getCertificateChain(alias)[0];
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} 
		return cert;
	}
	
	/**
	 * Gets a KeyPair for given alias 
	 * @param alias
	 * @param password
	 * @param storeFile
	 * @return
	 */
	public static KeyPair getKeyPair(String alias, char[] password, File storeFile) {
		KeyStore store;
		PublicKey pub = null;
		PrivateKey priv = null;
		FileInputStream fis;
		try {
			store = KeyStore.getInstance("IAIKKeyStore", "IAIK");
			fis = new FileInputStream(storeFile);
			store.load(fis, STORE_PWD.toCharArray());
			
//			Certificate cert = store.getCertificate(alias);
//			pub = cert.getPublicKey();
			KeyStore.PasswordProtection pwdProt = new KeyStore.PasswordProtection(password);
//			KeyStore.PrivateKeyEntry pkEntry = (KeyStore.PrivateKeyEntry) 
//				store.getEntry(alias, pwdProt);
//			priv = pkEntry.getPrivateKey();
	    priv = (PrivateKey)store.getKey(alias, password);
	    Certificate cert = store.getCertificateChain(alias)[0];
	    pub = cert.getPublicKey();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (UnrecoverableEntryException e) {
			e.printStackTrace();
		}
		
		return new KeyPair(pub,priv);
		
	}
	
	/**
	 * Outputs keystore contains for an alias. 
	 * 
	 * @author IAIK demo classes
	 * @param password
	 * @param fis
	 * @throws Exception
	 */
	public static void readKeyStore(char[] password, FileInputStream fis) throws Exception {
    // load keystore
    System.out.println("Load keystore...");
    KeyStore keyStore = KeyStore.getInstance("IAIKKeyStore", "IAIK");
    keyStore.load(fis, password);  
    // query aliases
    System.out.println("The keystore contains the following entries:");
    Enumeration aliases = keyStore.aliases();
    while (aliases.hasMoreElements()) {
      System.out.println(aliases.nextElement());   
    }    
    // fetch key
    PrivateKey privateKey = (PrivateKey)keyStore.getKey("Circus Authorities", password);
    System.out.println("Key is " + privateKey.getAlgorithm() + " (" + Util.getKeyLength(privateKey) + " bits)");
    // fetch certificates
    Certificate[] certChain = keyStore.getCertificateChain("Circus Authorities");
    System.out.println("Certificate chain is: ");
    for (int i = 0; i < certChain.length; i++) {
      System.out.println(((X509Certificate)certChain[i]).getSubjectDN());   
    }  
  }
	
	/**
	 * Saves given certificate chain to given keystore
	 * @param alias
	 * @param password
	 * @param pk
	 * @param cert
	 * @param storeFile
	 */
	public static void saveCert (String alias, char[] password, PrivateKey pk,
											 Certificate[] certs, File storeFile) {
		
		try {
			KeyStore store = KeyStore.getInstance("IAIKKeyStore", "IAIK");
			FileInputStream fis;
			fis = new FileInputStream(storeFile);
			store.load(fis, STORE_PWD.toCharArray());
	    store.setKeyEntry(alias, pk, password, certs);
//			store.setCertificateEntry(alias, cert);
			FileOutputStream fos = new FileOutputStream(storeFile);
			store.store(fos, STORE_PWD.toCharArray());
			fis.close();			
			fos.close();
		} catch (FileNotFoundException e1) {
			e1.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
		}

		Arrays.fill(password, '\u0000');
	}
	
	/**
	 * Generates 512-bit RSA KeyPair
	 * @return
	 */
	 public static KeyPair generateRSAKeyPair() {
	  	KeyPairGenerator generator=null;
			try {
				generator = KeyPairGenerator.getInstance("RSA", "IAIK");
			} catch (NoSuchAlgorithmException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (NoSuchProviderException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
			generator.initialize(512);
			return generator.generateKeyPair();
	  }
	 
	 public static void main(String []args){
			
		 try {
	      Security.insertProviderAt(new IAIK(), 2);
	      Security.insertProviderAt(new MDBProvider(), 3);
	    } catch( Throwable e ) {
	      //
	    }
		 
		 // KeyPair kp = generateRSAKeyPair();
			File file = new File(KEYSTORE);
			FileOutputStream fos=null;
			try {
				fos = new FileOutputStream(file);
			} catch (FileNotFoundException e) {
				e.printStackTrace();
			}
			try {
				
//				System.out.println(kp.getPrivate().toString());
//				createKeyStore("salasana".toCharArray(),fos);
//				saveCert("tonni","salasana".toCharArray(),kp.getPrivate(),null, file);
				//System.out.println(getPrivateKey("tonni", "salasana".toCharArray(), file).toString());
			} catch (Exception e) {
				e.printStackTrace();
			} 
			try {
				fos.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
	 }
	 
}


